import pandas as pd

import input


# 分析教练效应
def calculate_coach_effect(coach_name, country, years, sport, gender, type, athletes):
    # 筛选出该国家和该运动的数据
    country_athletes = athletes[(athletes["Team"] == country) & (athletes["Sport"] == sport) & (athletes["Sex"] == gender)]
    # 按年份排序
    country_athletes_sorted = country_athletes.sort_values(by="Year")
    # 计算每年的奖牌总数和金牌数
    medal_summary = country_athletes_sorted.groupby("Year").agg(
        Total_Medals=("Medal", lambda x: (x != "No medal").sum()),
        Gold_Medals=("Medal", lambda x: (x == "Gold").sum())
    ).reset_index()
    # 计算每年相对于上一年的奖牌总数和金牌数的增加量
    medal_summary['Total_Medals_Change'] = medal_summary['Total_Medals'].diff().fillna(0)
    medal_summary['Gold_Medals_Change'] = medal_summary['Gold_Medals'].diff().fillna(0)
    # 筛选教练存在年份的数据，排除上一年也有教练的年份
    medal_athlete = medal_summary[(medal_summary['Year'].isin(years)) & (~medal_summary['Year'].shift(1).isin(years))]
    # 计算教练存在年份的平均奖牌总数和金牌数的变化
    avg_total_medals_change = medal_athlete['Total_Medals_Change'].mean()
    avg_gold_medals_change = medal_athlete['Gold_Medals_Change'].mean()
    # 根据类型调整结果
    if type == "C": # 集体项目
        avg_total_medals_change = 1 if avg_total_medals_change > 1 else avg_total_medals_change
        avg_gold_medals_change = 1 if avg_gold_medals_change > 1 else avg_gold_medals_change
    elif type == "B": # 集体项目和个人项目都有
        avg_total_medals_change /= 2
        avg_gold_medals_change /= 2

    coach_effect_1 = avg_total_medals_change
    coach_effect_2 = avg_gold_medals_change

    return coach_effect_1, coach_effect_2



if __name__ == "__main__":
    # 设置文件路径
    file_path = "./data/"  # 替换为你的文件路径
    # 加载数据
    dataframes = input.load_data(file_path)
    athletes = dataframes["athletes"]

    # coaches_data 已经定义
    coaches_data = input.load_json("./json/coaches_data.json")

    # 计算每位教练的效应
    coach_effects = []
    for coach in coaches_data:
        coach_effect_1, coach_effect_2 = calculate_coach_effect(coach["Coach_Name"], coach["Country"], coach["Years"], coach["Sport"], coach["Gender"], coach["Type"], athletes)
        coach_effect_strength = coach_effect_1 * 0.4 + coach_effect_2 * 0.6
        coach_effects.append({
            "Coach_Name": coach["Coach_Name"],
            "Country": coach["Country"],
            "Sport": coach["Sport"],
            "Type" : coach["Type"],
            "Coach_Effect_Total": coach_effect_1,
            "Coach_Effect_Gold": coach_effect_2,
            "Coach_Effect_Strength": coach_effect_strength,
        })

    # 改变数据类型
    coach_effects_df = pd.DataFrame(coach_effects)
    # 输出结果
    print(coach_effects_df)
    coach_effects_df.to_csv("coach_effects_df.csv", index=False, encoding="utf-8")
    print("处理后的数据已保存到 coach_effects_df.csv 文件中。")